import time

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np

# TSP问题可视化函数
def tsp_visualization(ax):
    ax.clear()
    np.random.seed(int(time.time()))
    cities = np.random.rand(10, 2)  # 生成随机城市坐标
    num_cities = len(cities)
    start_city = 0  # 从第一个城市开始
    path = [start_city]  # 初始化路径，包含起始城市
    correct_path = []  # 存储正确的路径

    # 计算两个城市之间的距离
    def distance(a, b):
        return np.linalg.norm(cities[a] - cities[b])

    # 更新动画每一帧的函数
    def update(frame):
        nonlocal path, correct_path
        if len(path) < num_cities:
            current_city = path[-1]
            unvisited = [city for city in range(num_cities) if city not in path]
            if unvisited:
                next_city = min(unvisited, key=lambda city: distance(current_city, city))
                path.append(next_city)  # 添加下一个城市到路径

                # 判断当前路径是否正确
                if len(path) == num_cities:
                    path.append(start_city)  # 最后一个城市连接回起始城市形成闭环
                    total_distance = calculate_total_distance(path)
                    if len(correct_path) == 0 or total_distance < calculate_total_distance(correct_path):
                        correct_path = path.copy()
                    path.pop()  # 移除闭环的城市

        ax.clear()
        ax.plot(cities[:, 0], cities[:, 1], 'o')  # 绘制城市

        for i in range(len(path) - 1):
            city1 = path[i]
            city2 = path[i + 1]
            if path[-1] == start_city:
                if path[:-1] == correct_path:
                    ax.plot([cities[city1, 0], cities[city2, 0]], [cities[city1, 1], cities[city2, 1]], 'g-')
                else:
                    ax.plot([cities[city1, 0], cities[city2, 0]], [cities[city1, 1], cities[city2, 1]], 'r-')
            else:
                ax.plot([cities[city1, 0], cities[city2, 0]], [cities[city1, 1], cities[city2, 1]], 'k-')

        for i in range(len(correct_path) - 1):
            city1 = correct_path[i]
            city2 = correct_path[i + 1]
            ax.plot([cities[city1, 0], cities[city2, 0]], [cities[city1, 1], cities[city2, 1]], 'g-')

        # 显示起始城市
        ax.scatter(cities[start_city, 0], cities[start_city, 1], s=100, c='red', label='Start City')

        # 设置标题和图例
        ax.set_title("TSP Visualization")
        ax.legend()

    # 计算路径总距离
    def calculate_total_distance(path):
        total_distance = 0.0
        for i in range(len(path) - 1):
            total_distance += distance(path[i], path[i + 1])
        return total_distance

    # 创建动画
    ani = animation.FuncAnimation(ax.figure, update, frames=range(num_cities * 2), interval=1000, repeat=False)
    return ani
